import torch
from transformers import BertTokenizer, BertForSequenceClassification
from sklearn.preprocessing import LabelEncoder
import pandas as pd
import matplotlib.pyplot as plt


def topic_recognize():
    # 加载标签编码器
    label_encoder = LabelEncoder()
    # 用训练时的标签类别来拟合label_encoder
    train_data_path = './data_analysis/topic_recognize/web_text_zh_train.csv'
    train_data = pd.read_csv(train_data_path)[:200000]
    train_labels = train_data['label'].values
    label_encoder.fit(train_labels)

    # 指定保存模型和tokenizer路径
    saved_model_path = './data_analysis/topic_recognize/saved_model'

    # 加载BERT tokenizer和模型
    tokenizer = BertTokenizer.from_pretrained(saved_model_path)
    model = BertForSequenceClassification.from_pretrained(saved_model_path)

    # 检查是否有可用的GPU
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    print(f"Using device: {device}")

    # 预测函数
    def predict_topic(text):
        model.eval()
        with torch.no_grad():
            encoding = tokenizer(text, return_tensors="pt", padding=True, truncation=True).to(device)
            output = model(encoding['input_ids'], attention_mask=encoding['attention_mask'])
            prediction = torch.argmax(output.logits, dim=1).item()
            topic = label_encoder.inverse_transform([prediction])
            return topic[0]

    # 读取Excel文件
    file_path = "../public/data/tables/data_orin.xlsx"
    df = pd.read_excel(file_path)

    # 预测话题并新增列
    df['话题预测'] = df['title'].apply(predict_topic)

    # 找出“话题预测”列中重复次数最多的前十个数据
    top_10_topics = df['话题预测'].value_counts().head(10)

    # 设置字体和图表风格
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用SimHei字体以支持中文
    plt.rcParams['font.size'] = 14  # 设置字体大小
    plt.rcParams['axes.labelweight'] = 'bold'  # 设置坐标轴标签字体粗细
    plt.rcParams['axes.titlesize'] = 16  # 设置标题字体大小
    plt.rcParams['axes.titleweight'] = 'bold'  # 设置标题字体粗细
    plt.rcParams['xtick.labelsize'] = 12  # 设置x轴刻度字体大小
    plt.rcParams['ytick.labelsize'] = 12  # 设置y轴刻度字体大小

    # 生成颜色列表
    colors = ['#FF9999', '#FFCC99', '#99CC99', '#66CCCC', '#99CCFF', '#CC99FF', '#FF99CC', '#FF6666', '#FF9933', '#99CCFF']

    # 生成柱状图
    plt.figure(figsize=(12, 8))
    bars = plt.bar(top_10_topics.index, top_10_topics.values, width=0.9, color=colors)  # 设置不同颜色
    plt.title('前十个话题预测的数量', fontsize=18, fontweight='bold', color='#333333')  # 设置标题颜色
    plt.xlabel('话题预测', fontsize=16, fontweight='bold', color='#333333')  # 设置x轴标签颜色
    plt.ylabel('数量', fontsize=16, fontweight='bold', color='#333333')  # 设置y轴标签颜色
    plt.xticks(rotation=45, color='#333333')  # 设置x轴刻度标签颜色
    plt.yticks(color='#333333')  # 设置y轴刻度标签颜色

    # 显示柱子顶部的数值
    for bar in bars:
        yval = bar.get_height()
        plt.text(bar.get_x() + bar.get_width()/2 - 0.1, yval + 0.5, int(yval), va='bottom', fontsize=12, fontweight='bold', color='#333333')

    plt.savefig('../public/data/img/topic.png')
    # plt.show()

    # 生成爆炸效果
    explode = [0.1] + [0] * (len(top_10_topics) - 1)  # 仅爆炸第一个切片（即第一个话题）

    # 绘制饼图
    fig, ax = plt.subplots(figsize=(12, 8))
    ax.pie(top_10_topics, explode=explode, labels=top_10_topics.index, colors=colors, autopct='%1.1f%%',
           shadow=True, startangle=90, textprops={'fontsize': 12, 'fontweight': 'bold', 'color': '#333333'})
    ax.axis('equal')  # 确保饼图是一个正圆

    plt.title('前十个话题的比例', fontsize=18, fontweight='bold', color='#333333')
    plt.savefig('../public/data/img/topic_pie.png')
    # plt.show()
